import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tools.sm_exceptions import ConvergenceWarning
import warnings


# 数据处理函数（沿用之前的）
def process_data():
    excel_file = pd.ExcelFile("ori_message.xlsx")
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result['时间'] = pd.to_datetime(result['年'].astype(str) + '-' + result['月'].astype(str) + '-01')
    return result


if __name__ == "__main__":
    processed_data = process_data()
    processed_data.set_index('时间', inplace=True)  # 设置时间为索引
    flow_series = processed_data['流量(m3/s)']

    # 忽略收敛警告
    warnings.filterwarnings("ignore", category=ConvergenceWarning)

    # 定义AR(p)和MA(q)的阶数范围
    p_values = range(1, 6)  # p=1,2,3,4,5
    q_values = range(0, 6)  # q=0,1,2,3,4,5

    aic_matrix = []
    for p in p_values:
        row = []
        for q in q_values:
            try:
                # 对二阶差分后的数据进行建模（假设二阶差分后平稳）
                diff1 = flow_series.diff(1).dropna()
                diff2 = diff1.diff(1).dropna()
                stationary_series = diff2

                model = ARIMA(stationary_series, order=(p, 0, q))
                results = model.fit()
                row.append(results.aic)
            except:
                row.append(None)  # 处理拟合失败的情况
        aic_matrix.append(row)

    # 构建AIC值的DataFrame
    aic_df = pd.DataFrame(
        aic_matrix,
        index=[f'AR{p}' for p in p_values],
        columns=[f'MA{q}' for q in q_values]
    )

    # 绘制热力图
    plt.figure(figsize=(10, 8))
    sns.heatmap(
        aic_df,
        annot=True,
        cmap='YlGnBu',
        cbar_kws={'label': 'AIC'},
        fmt=".2f"
    )
    plt.title('图25.热力图')
    plt.xlabel('MA阶数')
    plt.ylabel('AR阶数')
    plt.show()